import os
import json
import cv2
import numpy as np
from matplotlib import pyplot as plt
from PIL import Image

data_dir = r'C:\Users\jiangzhengquan\Desktop\Switch_Dataset\JPEGImages'
save_dir = r'C:\Users\jiangzhengquan\Desktop\Switch_Dataset\SegmentationClass'

for file in os.listdir(data_dir):
    if file.endswith('.json'):
        name = os.path.splitext(file)[0]
        img_name = os.path.join(data_dir, name + '.jpg')
        content = json.load(open(os.path.join(data_dir, file), 'r'))
        img = cv2.imread(img_name)

        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        mask = np.zeros_like(img, dtype=np.uint8)
        for shape in content['shapes']:
            point = shape['points']
            label = shape['label']
            point = np.asarray(point, dtype=np.int32)[None]
            cv2.drawContours(mask, point, -1, 255, 2)
            cv2.drawContours(mask, point, -1, int(label), -1)
            # if img.shape[0] != 1080:
            #     img = cv2.resize(img, (1920, 1080))
        save_img = Image.fromarray(mask)
        save_img.save(os.path.join(save_dir, name + '.png'))
    print(file)
